###This is a vectorized version of the 1D Richards equation.
## The central difference scheme is used to approximate the spatial derivatives and the resulting
##ODEs are solved with a suitable integrator
import  numpy as np 
from parameters_mpc import spatialVariables_1D 
import van_Genuchten_relations as vg 

def Richards1D_vectorized(head, irrigAmount, cropCoeff, refEvap, soilPars):
    
    ##The spatial parameters
    totalDepth, axialNodes, axialDistance, totalNodes, numOfActuators, interPoints = spatialVariables_1D()
    capCapacity = vg.capillaryCapacity(head, soilPars) 

    ##Creating an array for the flux  at each node
    flux = np.zeros(totalNodes+1)

    ##Top boundary condition
    flux[totalNodes] = irrigAmount

    ##Bottom boundary
    #For the free drainage boundary condition, the flux equals the hydraulic conductivity of the bottom node
    flux[0] = -1*vg.hydraulicConductivity(head[0], soilPars)

    ##Flux for the internal nodes
    #The hydraulic conductivity at the center of the compartments
    cntr_1 = np.arange(0, totalNodes-1)
    nodalHydCond = vg.hydraulicConductivity(head, soilPars)
    approxHydCond = (nodalHydCond[cntr_1 + 1] + nodalHydCond[cntr_1])/2.0

    cntr_2 = np.arange(1,totalNodes)

    flux[cntr_2] = -approxHydCond*(((head[cntr_1 + 1]-head[cntr_1])/axialDistance) + 1.0)

    sinkTerms = np.zeros((1, axialNodes))

    
    i= np.arange(0, 1)
    distance = axialDistance*np.arange(0, axialNodes)
    sinkTerms[i,:] = (2*(refEvap * cropCoeff)/ totalDepth)*(1-(totalDepth - distance)/totalDepth) #Weather conditions
    
    
    cntr_3 = np.arange(0, totalNodes)
    
    dhdt = ((-(flux[cntr_3 + 1] - flux[cntr_3])/ axialDistance) - sinkTerms)/capCapacity
    return dhdt


def volMoistureAllNodes(head, soilPars):
    totalDepth, axialNodes, axialDistance, totalNodes, numOfActuators, interPoints = spatialVariables_1D()
    volMoistureInit = vg.volumetricMoisture(head, soilPars)
    volMoistureFin =volMoistureInit
    return volMoistureFin

